import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2
from segment_anything.utils.transforms import ResizeLongestSide
import sys
import os
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
def show_mask(mask, ax, random_color=True):
    """
    显示遮罩图像。

    参数:
    mask: 遮罩图像。
    ax: Matplotlib的子图对象。
    random_color: 是否使用随机颜色显示遮罩。
    """
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30/255, 144/255, 255/255, 0.6])
    h, w = mask.shape[-2:]
    print("h,w2:")
    print(h, w)
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)
def show_mask_patch(mask, ax, random_color=True):
    """
    显示遮罩覆盖的图像，遮罩未覆盖的图像则覆盖黑色处理

    参数:
    mask: 遮罩图像。
    ax: Matplotlib的子图对象。
    random_color: 是否使用随机颜色显示遮罩。
    """
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([0, 0, 0, 0.6])
    h, w = mask.shape[-2:]
    mask_2d = mask.reshape((h, w))
    mask_image = np.zeros((h, w, 4))  # 创建一个新的四通道图像（RGBA）
    mask_image[mask_2d == 1] = color  # 只对遮罩的前景部分应用颜色
    ax.imshow(mask_image)

def get_bbox_from_mask(mask):
    # 找到所有非零像素的坐标
    nonzero = np.nonzero(mask)
    if len(nonzero) == 2:
        rows, cols = nonzero
    elif len(nonzero) == 3:
        slices, rows, cols = nonzero
    else:
        raise ValueError("Unexpected number of dimensions in mask")
    # 获取边界框的坐标
    xmin = np.min(cols)
    ymin = np.min(rows)
    xmax = np.max(cols)
    ymax = np.max(rows)
    return xmin, ymin, xmax, ymax
def show_points(coords, labels, ax, marker_size=375):
    """
    显示点。

    参数:
    coords: 点的坐标。
    labels: 点的标签。
    ax: Matplotlib的子图对象。
    marker_size: 点的大小。
    """
    if coords is None or labels is None:
        return
    pos_points = coords[labels==1]
    neg_points = coords[labels==0]
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*',
               s=marker_size, edgecolor='white', linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*',
               s=marker_size, edgecolor='white', linewidth=1.25)
#显示标记点（鼠标）
def show_anns(anns):
    if len(anns) == 0:
        return
    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
    ax = plt.gca()
    ax.set_autoscale_on(False)

    img = np.ones((sorted_anns[0]['segmentation'].shape[0], sorted_anns[0]['segmentation'].shape[1], 4))
    img[:,:,3] = 0
    for ann in sorted_anns:
        m = ann['segmentation']
        color_mask = np.concatenate([np.random.random(3), [0.35]])
        img[m] = color_mask
    ax.imshow(img)
def show_box(box, ax):
    """
    显示矩形框。

    参数:
    box: 矩形框的坐标。
    ax: Matplotlib的子图对象。
    """
    # 如果box为空，初始化它为一个空列表
    if box is None:
        box = [0, 0, 0, 0]

    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0),
                               lw=2))

def load_image(image_path):
    """
    加载图像,并将其转换为模型课接受图像。

    参数:
    image_path: 图像的路径。

    返回:
    image: 加载的图像。
    """
    # image = cv2.imread(image_path)
    # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # return image
    # Read image with OpenCV
    image = cv2.imread(image_path)

    # Check if image is empty
    if image is None:
        print(f"Error: Unable to read image from {image_path}")
        return None

    # Convert to RGB
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    return image

def load_sam_model():
    """
    加载SAM模型。

    参数:
    checkpoint_path: 模型检查点的路径。
    model_type: 模型类型。
    device: 设备类型。

    返回:
    sam: 加载的SAM模型。
    predictor: SAM模型的预测器。
    """
    # 写死的参数
    sam_checkpoint = os.path.join("service","segment","sam_vit_h_4b8939.pth")
    model_type = "vit_h"
    sys.path.append("..")
    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
    device = "cuda:1" if torch.cuda.is_available() else "cpu"
    sam.to(device=device)
    predictor = SamPredictor(sam)
    return sam, predictor

def process_image(predictor, image):
    """
    处理图像,通过调用SamPredictor来处理图像以生成图像嵌入。设置图像。SamPredictor会记住这个嵌入，并将其用于后续的掩码预测。。

    参数:
    predictor: SAM模型的预测器。
    image: 要处理的图像。
    """
    predictor.set_image(image)

def predict_mask(predictor, input_point, input_label, multimask_output, mask_input, input_box):
    """
    预测遮罩。

    参数:
    predictor: SAM模型的预测器。
    input_point: 输入点的坐标。
    input_label: 输入点的标签。
    multimask_output: 是否输出多个遮罩。
    mask_input: 输入的遮罩。
    input_box: 输入的矩形框。

    返回:
    masks: 预测的遮罩。
    scores: 遮罩的分数。
    logits: 遮罩的logits。
    """
    if mask_input is not None:
        mask_input = mask_input[None, :, :]
    else:
        mask_input = None

    if input_box is not None:
        box = input_box[None, :]
    else:
        box = None

    masks, scores, logits = predictor.predict(
        point_coords=input_point,
        point_labels=input_label,
        mask_input=mask_input,
        box=box,
        multimask_output=multimask_output,
    )
    return masks, scores, logits
#多框遮罩预测
def predict_mask2(image,predictor, multimask_output, input_boxes):
    """
    预测遮罩。

    参数:
    predictor: SAM模型的预测器。
    input_point: 输入点的坐标。
    input_label: 输入点的标签。
    multimask_output: 是否输出多个遮罩。
    mask_input: 输入的遮罩。
    input_boxes: 输入的矩形框。

    返回:
    masks: 预测的遮罩。
    scores: 遮罩的分数。
    logits: 遮罩的logits。
    """

    input_boxes = input_boxes.to(predictor.device)

    transformed_boxes = predictor.transform.apply_boxes_torch(input_boxes, image.shape[:2])
    masks, _, _ = predictor.predict_torch(
        point_coords=None,
        point_labels=None,
        boxes=transformed_boxes,
        multimask_output= multimask_output,
    )
    return masks
def generate_image_number(directory):
    """
    生成新的图片编号。

    参数:
    directory: 图片保存的目录。

    返回:
    new_image_number: 新的图片编号。
    """
    # 获取目录下的文件数量
    file_count = len(os.listdir(directory))
    # 生成新的图片编号
    new_image_number = str(file_count + 1).zfill(8)
    return new_image_number

#新图片生成函数--适配单点，多点，点和矩形框，矩形框
def Simple_coordinate_new_image( images, points, labels, boxes, masks):
    """
    根据遮罩图像和原图像生成新的图像。

    参数:
    predictor: SAM模型的预测器。
    images: 图像列表。
    points: 点的坐标列表。
    labels: 点的标签列表。
    boxes: 矩形框列表。
    multimask_output: 是否输出多个遮罩。
    masks: 预测的遮罩列表。

    返回:
    save_paths1：保存的原图像和遮罩图像的路径列表-后端用。
    save_paths2：保存的原图像,遮罩和标记点的路径列表-前端展示用。
    """
    save_paths1 = []
    save_paths2 = []
    for image, input_point, input_label, input_box, mask in zip(images, points, labels, boxes, masks):
        plt.figure(figsize=(10, 10))
        plt.imshow(image)
        show_mask(mask[0], plt.gca())
        plt.axis('off')
        #在这个位置保不存在标记点的图片
        save_dir1 = 'SaveTheImage'
        # 获取新的图片编号
        image_number = generate_image_number(save_dir1)
        if not os.path.exists(save_dir1):
            os.makedirs(save_dir1)
        # 生成保存图像的路径
        save_path1 = os.path.join(save_dir1, f'{image_number}.png')
        plt.savefig(save_path1)  # 保存图像到指定路径
        print("保存图像到指定路径：" + save_path1)

        #绘制标记点
        show_points(input_point, input_label, plt.gca())
        if input_box is not None:
            show_box(input_box, plt.gca())  # 在图像上显示矩形框
        # 写死的保存图像的目录路径
        save_dir2 = 'DisplayPicture'
        if not os.path.exists(save_dir2):
            os.makedirs(save_dir2)
        # 生成保存图像的路径
        save_path2 = os.path.join(save_dir2, f'{image_number}.png')
        plt.savefig(save_path2)  # 保存图像到指定路径
        print("保存图像到指定路径：" + save_path2)
        plt.show() # 显示原图像和遮罩图像以及标记点
        plt.close()  # 关闭图像，防止内存泄漏

        save_paths1.append(save_path1)
        save_paths2.append(save_path2)

    return save_paths1, save_paths2
def Simple_coordinate_new_image2 ( images, boxes, masks):
    """
    根据遮罩图像和原图像生成新的图像。

    参数:
    predictor: SAM模型的预测器。
    images: 图像列表。
    boxes: 矩形框列表。
    masks: 预测的遮罩列表。

    返回:
    save_paths1：保存的原图像和遮罩图像的路径列表-后端用。
    save_paths2：保存的原图像,遮罩和标记点的路径列表-前端展示用。
    """
    save_paths1 = []
    save_paths2 = []
    for image,  input_boxes, mask in zip(images,  boxes, masks):
        plt.figure(figsize=(10, 10))
        plt.imshow(image)
       #绘制遮罩
        for(mask_index,mask) in enumerate(mask):
            show_mask(mask.cpu().numpy(), plt.gca(), random_color=True)

        plt.axis('off')
        #在这个位置保不存在标记点的图片
        save_dir1 = 'SaveTheImage'
        # 获取新的图片编号
        image_number = generate_image_number(save_dir1)
        if not os.path.exists(save_dir1):
            os.makedirs(save_dir1)
        # 生成保存图像的路径
        save_path1 = os.path.join(save_dir1, f'{image_number}.png')
        plt.savefig(save_path1)  # 保存图像到指定路径
        print("保存图像到指定路径：" + save_path1)

        # 绘制框
        #if input_boxes is not None:
        for box in input_boxes:
            show_box(box.cpu().numpy(), plt.gca())
        # 写死的保存图像的目录路径
        save_dir2 = 'DisplayPicture'
        if not os.path.exists(save_dir2):
            os.makedirs(save_dir2)
        # 生成保存图像的路径
        save_path2 = os.path.join(save_dir2, f'{image_number}.png')
        plt.savefig(save_path2)  # 保存图像到指定路径
        print("保存图像到指定路径：" + save_path2)
        plt.show() # 显示原图像和遮罩图像以及标记点
        plt.close()  # 关闭图像，防止内存泄漏

        save_paths1.append(save_path1)
        save_paths2.append(save_path2)
        return save_paths1, save_paths2


# 生成自动分割图像
def Simple_coordinate_new_image3(image, mask2):
    """
    根据遮罩图像和原图像生成新的图像。

    参数:。
    images: 图像。
    mask: 预测的遮罩。

    返回:
    save_paths：保存的原图像和遮罩图像-前后端用。
    """
    save_paths = []
    plt.figure(figsize=(10, 10))
    plt.imshow(image)
    show_anns(mask2)
    plt.axis('off')
    # 在这个位置保不存在标记点的图片
    save_dir = 'SaveTheImage'
    # 获取新的图片编号
    image_number = generate_image_number(save_dir)
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    # 生成保存图像的路径
    save_path = os.path.join(save_dir, f'{image_number}.png')
    plt.savefig(save_path)  # 保存图像到指定路径
    print("保存图像到指定路径：" + save_path)
    plt.show()  # 显示原图像和遮罩图像以及标记点
    plt.close()  # 关闭图像，防止内存泄漏
    save_paths.append(save_path)
    return save_paths



# 批量处理相关
def get_resize_transform(sam):
    """
    获取调整图像大小的转换。

    参数:
    sam: SAM模型。

    返回:
    resize_transform: 调整图像大小的转换。
    """
    resize_transform = ResizeLongestSide(sam.image_encoder.img_size)
    return resize_transform


#进行图片预处理，后将处理后的图像转换为PyTorch张量，并确保它在正确的设备上。
def prepare_image(image, transform, device):
    """
    准备图像以供模型使用。

    参数:
    image: 原始图像。
    transform: 用于预处理图像的转换对象。
    device: 用于执行计算的设备（CPU或GPU）。

    返回:
    返回预处理后的图像。
    """
    image = transform.apply_image(image)
    image = torch.as_tensor(image, device=device.device)
    return image.permute(2, 0, 1).contiguous()

def get_batched_input(images, boxes, model,resize_transform):
    """
        获取批量输入。

        参数:
        images: 图像列表。
        boxes: 边界框列表。
        transform: 用于预处理图像的转换对象。
        model: 使用的模型。
        resize_transform: 用于调整图像大小的转换对象。

        返回:
        返回批量输入的列表，每个输入包含图像、边界框和原始大小。
        """
    batched_input = []
    for image, box in zip(images, boxes):
        input_dict = {
            'image': prepare_image(image, resize_transform, model),
            'boxes': resize_transform.apply_boxes_torch(box, image.shape[:2]),
            'original_size': image.shape[:2]
        }
        batched_input.append(input_dict)
        print(batched_input)
    return batched_input

def batch_process_images(batched_output,boxes, images):
    """
    批量处理图像。

    参数:
    batched_output: 包含预测遮罩和可能的边界框的字典列表。
    images: 要处理的图像列表。

    对于每个图像和对应的输出字典，提取遮罩和边界框（如果存在），然后调用
    Simple_coordinate_new_image2函数生成新的图像。
    """

    for output_dict,boxes,image in zip(batched_output,boxes,images):
        masks = output_dict['masks']
        boxes = output_dict.get('boxes', boxes)
        Simple_coordinate_new_image2([image], [boxes], [masks])



#自动分割
def generate_masks(image, model, points_per_side=32, pred_iou_thresh=0.86, stability_score_thresh=0.92, crop_n_layers=1, crop_n_points_downscale_factor=2, min_mask_region_area=100):
    #from segment_anything import SamAutomaticMaskGenerator

    mask_generator = SamAutomaticMaskGenerator(
        model=model,
        points_per_side=points_per_side,
        pred_iou_thresh=pred_iou_thresh,
        stability_score_thresh=stability_score_thresh,
        crop_n_layers=crop_n_layers,
        crop_n_points_downscale_factor=crop_n_points_downscale_factor,
        min_mask_region_area=min_mask_region_area
    )

    masks = mask_generator.generate(image)

    return masks

#测试：
if __name__ == '__main__':
    image2 = load_image("image/truck.jpg")
    image1 = load_image("image/groceries.jpg")
    sam, predictor = load_sam_model()
    process_image(predictor, image1)
    process_image(predictor, image2)
    # 是否输出多个遮罩
    multimask_output = False
    # 分批提示输入
    input_boxes2 = torch.tensor([
        [75, 275, 1725, 850],
        [425, 600, 700, 875],
        [1375, 550, 1650, 800],
        [1240, 675, 1400, 750],
    ], device=predictor.device)
    input_boxes1 = torch.tensor([
        [450, 170, 520, 350],
        [350, 190, 450, 350],
        [500, 170, 580, 350],
        [580, 170, 640, 350],
    ], device=predictor.device)
    # 调整图像大小的转换
    resize_transform = get_resize_transform(sam)
    # 生成批量输入
    batched_input = get_batched_input([image1, image2], [input_boxes1, input_boxes2], sam, resize_transform)
    # 批量预测
    batched_output = sam(batched_input, multimask_output=False)
    # 批量处理
    batch_process_images(batched_output, [input_boxes1, input_boxes2], [image1, image2])

